import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import find_peaks, savgol_filter

def _generate_simulation_signal(sampling_rate=1000, duration=5):
    # 生成一个模拟信号，大多数时候是一条直线，偶尔有几个波
    t = np.linspace(0, duration, int(sampling_rate * duration), endpoint=False)
    
    # 初始化为一条直线
    signal = np.zeros(len(t))
    
    # 在死区外添加单波
    outside_dead_zone_start = int(0.2 * sampling_rate)
    outside_dead_zone_end = int(0.4 * sampling_rate)
    # 计算单波的频率，使得在指定时间段内只有一个完整周期
    single_wave_freq = 1 / ((outside_dead_zone_end - outside_dead_zone_start) / sampling_rate)
    signal[outside_dead_zone_start:outside_dead_zone_end] += 0.5 * np.sin(2 * np.pi * single_wave_freq * t[outside_dead_zone_start:outside_dead_zone_end])
    
    # 修改为整数索引
    start_idx_1 = int(4 * sampling_rate)
    end_idx_1 = int(4.5 * sampling_rate)
    start_idx_2 = int(2 * sampling_rate)
    end_idx_2 = int(3 * sampling_rate)
    start_idx_3 = int(3 * sampling_rate)
    end_idx_3 = int(3.3 * sampling_rate)
    start_idx_4 = int(1.5 * sampling_rate)
    end_idx_4 = int(1.8 * sampling_rate)

    # 添加一些幅度较小的“死区”内的单波
    freq_1 = 1 / ((end_idx_1 - start_idx_1) / sampling_rate)
    signal[start_idx_1 : end_idx_1] += 0.05 * np.sin(2 * np.pi * freq_1 * t[start_idx_1 : end_idx_1])
    
    freq_2 = 1 / ((end_idx_2 - start_idx_2) / sampling_rate)
    signal[start_idx_2 : end_idx_2] += 0.03 * np.sin(2 * np.pi * freq_2 * t[start_idx_2 : end_idx_2])
   
    freq_3 = 1 / ((end_idx_3 - start_idx_3) / sampling_rate)
    signal[start_idx_3 : end_idx_3] += 0.03 * np.sin(2 * np.pi * freq_3 * t[start_idx_3 : end_idx_3])
   
    freq_4 = 1 / ((end_idx_4 - start_idx_4) / sampling_rate)
    signal[start_idx_4 : end_idx_4] += 0.13 * np.sin(2 * np.pi * freq_4 * t[start_idx_4 : end_idx_4])

    return signal, t
def extract_waves_by_gradient_features(signal, 
                                       peak_prominence=0.1, 
                                       trough_prominence=0.1,
                                       min_peak_distance_samples=10):
    # 计算信号的均值作为死区阈值
    threshold = np.mean(signal)

    """
    根据梯度特征（平坦点 -> 波峰 -> 平坦点）提取定义的一个个完整波形。
    这里简化为识别波谷 -> 波峰 -> 波谷的波段。

    参数:
    signal (np.array): 一维输入信号。
    peak_prominence (float): 识别波峰的最小突出度。
    trough_prominence (float): 识别波谷的最小突出度。
    min_peak_distance_samples (int): 两个波峰/波谷之间最小的样本距离。

    返回:
    list: 包含每个提取波段信号的列表。
    list: 包含每个波段起始和结束索引的列表。
    """

    # 移除平滑信号和梯度计算部分

    # 2. 检测波峰和波谷
    # find_peaks 用于检测波峰，要检测波谷，可以对信号取反
    peaks_indices, _ = find_peaks(signal,
                                  prominence=peak_prominence * np.std(signal), # 使用信号标准差的倍数作为突出度
                                  distance=min_peak_distance_samples)

    # 找到波谷：对信号取反，再找波峰
    troughs_indices, _ = find_peaks(-signal,
                                    prominence=trough_prominence * np.std(signal),
                                    distance=min_peak_distance_samples)
    # 过滤掉不在阈值范围内的极值点
    peaks_indices = peaks_indices[np.abs(signal[peaks_indices] - threshold) < np.std(signal)]
    troughs_indices = troughs_indices[np.abs(signal[troughs_indices] - threshold) < np.std(signal)]
    #print("过滤后的波峰索引:", peaks_indices)
    #print("过滤后的波谷索引:", troughs_indices)
    
    # 合并并排序所有的极值点（波峰和波谷），作为潜在的波形边界
    all_extrema_indices = np.sort(np.unique(np.concatenate((peaks_indices, troughs_indices))))

    extracted_waves = []
    wave_indices = []

    # 3. 构建波段并提取
    # 我们定义一个波从一个波谷开始，经过一个波峰，到下一个波谷结束
    # 或者从一个梯度接近0的点（极值点）开始，到下一个梯度接近0的点结束

    # 计算信号梯度（使用简单差分近似）
    gradient = np.gradient(signal)
    # 定义梯度接近0的阈值（可调整）
    gradient_threshold = 0.01 * np.max(np.abs(gradient))
    # 计算信号的中位数
    median_value = np.median(signal)
    value_threshold = 0.01 * np.std(signal)
    # 检测梯度接近0且值在中位数附近的平坦点
    flat_indices = np.where((np.abs(gradient) < gradient_threshold) & (np.abs(signal - median_value) < value_threshold))[0]
    
    # 过滤掉前后梯度都为 0 的平坦点
    valid_flat_indices = []
    for idx in flat_indices:
        # 处理边界情况
        prev_idx = max(0, idx - 1)
        next_idx = min(len(gradient) - 1, idx + 1)
        
        if np.abs(gradient[prev_idx]) >= gradient_threshold or np.abs(gradient[next_idx]) >= gradient_threshold:
            valid_flat_indices.append(idx)
    
    flat_indices = np.array(valid_flat_indices)
    #print("过滤后的平坦索引:", flat_indices)
    
    # 合并所有特征点（平坦点、波峰、波谷）并排序
    all_feature_indices = np.unique(np.concatenate((flat_indices, peaks_indices, troughs_indices)))
    all_feature_indices.sort()
    extracted_waves = []
    wave_indices = []

    # 寻找有效的"平坦 -> 多个波峰波谷 -> 平坦"序列
    #print("所有特征点索引:", all_feature_indices, "特征点数量:", len(all_feature_indices))
    i = 0
    while i < len(all_feature_indices) - 1:
        # 找到起始的平坦点
        if all_feature_indices[i] in flat_indices:
            start_flat_idx = all_feature_indices[i]
            j = i + 1
            found_peak_trough = False
            while j < len(all_feature_indices):
                # 检查是否有波峰和波谷交替出现
                if all_feature_indices[j] in peaks_indices:
                    is_peak = True
                    j += 1
                    while j < len(all_feature_indices) and all_feature_indices[j] not in troughs_indices:
                        j += 1
                    if j < len(all_feature_indices):
                        is_trough = True
                        found_peak_trough = True
                        j += 1
                else:
                    break
            # 找到结束的平坦点
            if j < len(all_feature_indices) and all_feature_indices[j] in flat_indices:
                end_flat_idx = all_feature_indices[j]
                if end_flat_idx - start_flat_idx >= min_peak_distance_samples and found_peak_trough:
                    extracted_waves.append(signal[start_flat_idx:end_flat_idx+1])
                    wave_indices.append((start_flat_idx, end_flat_idx))
                    #print("提取的波形索引:", start_flat_idx, end_flat_idx)
                    i = j
                else:
                    i += 1
            else:
                i += 1
        else:
            i += 1

    # 4. 输出提取的波段数量
    return extracted_waves, wave_indices

def _plot_extracted_waves(t, signal, extracted_waves, wave_indices):
    """
    绘制原始信号和提取出的波形。

    参数:
    t (np.ndarray): 时间数组，用于表示信号的时间轴。
    signal (np.ndarray): 原始输入信号，一维数组。
    extracted_waves (list): 包含每个提取波段信号的列表。
    wave_indices (list): 包含每个提取波段起始和结束索引的列表。

    返回:
    无，直接显示绘制的图形并打印提取的波形数量。
    """
    # 3. 绘图展示结果
    plt.figure(figsize=(15, 8))

    plt.subplot(2, 1, 1)
    plt.plot(t, signal, label='Original Signal', alpha=0.7)
    plt.title('Original Signal')
    plt.xlabel('Time (s)')
    plt.ylabel('Amplitude')
    plt.legend()
    plt.grid(True)

    plt.subplot(2, 1, 2)
    plt.plot(t, signal, label='Original Signal', alpha=0.5, color='gray') # 背景显示原始信号

    colors = plt.cm.jet(np.linspace(0, 1, len(extracted_waves))) # 为每个波段生成不同颜色
    for i, (start_idx, end_idx) in enumerate(wave_indices):
        t_segment = t[start_idx:end_idx+1]
        signal_segment = signal[start_idx:end_idx+1]
        plt.plot(t_segment, signal_segment, color=colors[i], linewidth=2, label=f'Extracted Wave {i+1}' if i == 0 else "")
        plt.axvline(t[start_idx], color=colors[i], linestyle=':', linewidth=0.8)
        plt.axvline(t[end_idx], color=colors[i], linestyle=':', linewidth=0.8)

    plt.title('Extracted Waveforms')
    plt.xlabel('Time (s)')
    plt.ylabel('Amplitude')
    plt.legend()
    plt.grid(True)

    plt.tight_layout()
    plt.show()

    print(f"Total extracted waves: {len(extracted_waves)}")


# --- 示例用法 ---
if __name__ == "__main__":
    # 1. 生成一个模拟信号，大多数时候是一条直线，偶尔有几个波
    sampling_rate = 1000  # 1000 Hz
    duration = 5        # 5 秒
    signal, t = _generate_simulation_signal(sampling_rate, duration)
    
    # 2. 设置提取参数
    # peak_prominence 和 trough_prominence 是关键，它们决定了我们认为多大的“山峰”或“山谷”才算一个真正的波峰/波谷
    # 它们可以根据信号的噪声水平和波形特征进行调整
    # min_peak_distance_samples 避免检测到过密的“波峰/波谷”，防止噪声引起的误判
    extracted_waves, wave_indices= extract_waves_by_gradient_features(
        signal,
        peak_prominence=0.1,     # 峰值突出度百分比 (相对于信号标准差)
        trough_prominence=0.1,     # 波谷突出度百分比
        min_peak_distance_samples=int(sampling_rate / 200) # 最小波峰/谷距离 (例如20Hz波的半个周期)
    )

    # 打印提取的每个波段的起始和结束索引
    for i, index in enumerate(wave_indices, start=1):
        print(f"Extracted wave_indices {i}:")
        print(index)
        print("-" * 50)

    # 3. 绘图展示结果
    _plot_extracted_waves(t, signal, extracted_waves, wave_indices)
